from api import sqlpyodbc, sqldf, sql_polars
from pydantic import BaseModel
from typing import Optional
from datetime import datetime
import traceback

class StorageRecord(BaseModel):
    storage_type: str  # 'I' for in, 'O' for out
    storage_number: Optional[str] = None  # Required for 'I', optional for 'O'
    material_code: str
    quantity: int
    remark: Optional[str] = None
    operator: Optional[str] = None

    class Config:
        orm_mode = True

def get_storage_records(
    search_text: str = '', 
    storage_type: Optional[str] = None,
    start_date: Optional[str] = None, 
    end_date: Optional[str] = None, 
    page: int = 1, 
    page_size: int = 10
):
    # 首先获取总记录数
    count_sql = """
        SELECT COUNT(*) as total
        FROM StorageRecord s
        LEFT JOIN PackageMaterial p ON s.MaterialCode = p.Code
    """
    conditions = []
    
    if search_text:
        conditions.append(f"""(
            s.StorageNumber LIKE '%{search_text}%' 
            OR s.MaterialCode LIKE '%{search_text}%'
            OR p.Type LIKE '%{search_text}%'
            OR p.Description LIKE '%{search_text}%'
        )""")
    
    if storage_type:
        conditions.append(f"s.StorageType = '{storage_type}'")
    
    if start_date and end_date:
        # 处理日期格式，移除 T 和 Z
        start_date = start_date.split('T')[0] if 'T' in start_date else start_date
        end_date = end_date.split('T')[0] if 'T' in end_date else end_date
        conditions.append(
            f"CONVERT(date, s.CreateTime) >= '{start_date}' AND CONVERT(date, s.CreateTime) <= '{end_date}'"
        )
    
    if conditions:
        count_sql += " WHERE " + " AND ".join(conditions)
    
    try:
        total_result = sql_polars.query(count_sql)
        
        # 改进错误处理
        if total_result is None:
            print("Count query returned None")
            return {"records": [], "total": 0}
            
        if isinstance(total_result, bool):
            print("Count query returned boolean value")
            return {"records": [], "total": 0}
            
        total = int(total_result.select("total")[0, 0]) if total_result.height > 0 else 0
        
        if total == 0:
            return {"records": [], "total": 0}
        
        # 计算分页
        offset = (page - 1) * page_size
        
        # 获取分页数据
        sql = """
            SELECT s.StorageNumber as storageNumber,
                   s.MaterialCode as code,
                   s.Quantity as quantity,
                   s.StorageType as storageType,
                   s.Remark as remark,
                   s.CreateBy as operator,
                   s.CreateTime as createTime,
                   s.Accumulation as accumulation,
                   p.Type as type,
                   p.Description as description,
                   p.Unit as unit
            FROM StorageRecord s
            LEFT JOIN PackageMaterial p ON s.MaterialCode = p.Code
        """
        
        if conditions:
            sql += " WHERE " + " AND ".join(conditions)
        
        sql += f" ORDER BY s.CreateTime DESC OFFSET {offset} ROWS FETCH NEXT {page_size} ROWS ONLY"

        result = sql_polars.query(sql)
        
        # 改进错误处理
        if result is None or isinstance(result, bool):
            return {"records": [], "total": total}
            
        records = result.to_dicts()
        
        return {
            "records": records,
            "total": total
        }
    except Exception as e:
        print(f"Error in get_storage_records: {str(e)}")
        print(f"Error details: {traceback.format_exc()}")
        return {"records": [], "total": 0}

def add_storage_record(storage: StorageRecord):
    try:
        # 检查包材是否存在
        check_sql = f"""
            SELECT CurrentStock 
            FROM PackageMaterial 
            WHERE Code = '{storage.material_code}' 
            AND IsDeleted = 0
        """
        result = sqlpyodbc.Query_one(check_sql)
        
        if not result:
            return {
                "code": 400,
                "message": f"包材不存在: {storage.material_code}",
                "data": None
            }
            
        current_stock = result[0]
        # 如果是出库，检查库存是否充足
        if storage.storage_type == 'O' and current_stock < storage.quantity:
            return {
                "code": 400,
                "message": f"库存不足: 当前库存={current_stock}, 需要数量={storage.quantity}",
                "data": None
            }

        # 执行入库或出库操作
        sql = f"""
            INSERT INTO StorageRecord (
                StorageNumber,
                MaterialCode,
                Quantity,
                StorageType,
                Remark,
                CreateBy
            ) VALUES (
                {f"'{storage.storage_number}'" if storage.storage_number else 'NULL'},
                '{storage.material_code}',
                {storage.quantity},
                '{storage.storage_type}',
                '{storage.remark or ""}',
                '{storage.operator or "system"}'
            )
        """
        result = sqlpyodbc.execute(sql)
        if result:
            return {
                "code": 200,
                "message": "操作成功",
                "data": None
            }
        else:
            return {
                "code": 500,
                "message": "操作失败",
                "data": None
            }
    except Exception as e:
        print(f"Error in add_storage_record: {str(e)}")
        return {
            "code": 500,
            "message": f"操作失败: {str(e)}",
            "data": None
        }

def get_material_list():
    try:
        sql = """
            SELECT 
                Code,
                Type,
                Description,
                ISNULL(CurrentStock, 0) as CurrentStock
            FROM PackageMaterial 
            WHERE IsDeleted = 0
            ORDER BY Code
        """
        
        df = sql_polars.query(sql)
        
        if df is None or df.height == 0:
            return None

        df.columns = [col.lower() for col in df.columns]

        return df.to_dicts()
        
    except Exception as e:
        print(f"Error in get_material_list: {str(e)}")
        return None 